/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parvec_axpby.h"
#include "parvec_axpbypcz.h"
#include "par_compound_operator.h"

namespace YHAMG
{

Par_Compound_Operator::Par_Compound_Operator(double _beta, const Par_Operator& _C, double _gamma)
	: Par_Operator(_C.comm),
	A(0),
	B(0),
	C(&_C),
	alpha(0.0),
	beta(_beta),
	gamma(_gamma)
{
}

Par_Compound_Operator::Par_Compound_Operator(double _alpha, const Par_Operator& _B, double _beta, const Par_Operator& _C, double _gamma)
	: Par_Operator(_B.comm),
	A(0),
	B(&_B),
	C(&_C),
	alpha(_alpha),
	beta(_beta),
	gamma(_gamma)
{
}

Par_Compound_Operator::Par_Compound_Operator(double _alpha, const Par_Operator& _A, const Par_Operator& _B, double _gamma)
	: Par_Operator(_A.comm),
	A(&_A),
	B(&_B),
	C(0),
	alpha(_alpha),
	beta(0.0),
	gamma(_gamma)
{
}

Par_Compound_Operator::Par_Compound_Operator(double _alpha, const Par_Operator& _A, const Par_Operator& _B, double _beta, const Par_Operator& _C, double _gamma)
	: Par_Operator(_A.comm),
	A(&_A),
	B(&_B),
	C(&_C),
	alpha(_alpha),
	beta(_beta),
	gamma(_gamma)
{
}

int Par_Compound_Operator::InSize() const
{
	if (C) return C->InSize();
	if (B) return B->InSize();
	return 0;
}

int Par_Compound_Operator::OutSize() const
{
	if (C) return C->OutSize();
	if (A) return A->OutSize();
	if (B) return B->OutSize();
	return 0;
}

void Par_Compound_Operator::Apply(const Par_Vector& x, const Par_Vector& y) const
{
	if (alpha == 0 || !B)
	{
		if (beta == 0 || !C)
			ParVec_AXPBY(0.0, y, gamma, x, y);
		else
		{
			C->Apply(x, y);
			ParVec_AXPBY(beta, y, gamma, x, y);
		}
	}
	else
	{
		Par_Vector z(comm);
	
		if (A)
		{
			Par_Vector w(comm);
			w.Resize(B->OutSize());
			z.Resize(A->OutSize());
			B->Apply(x, w);
			A->Apply(w, z);
		}
		else
		{
			z.Resize(B->OutSize());
			B->Apply(x, z);
		}
		
		if (beta == 0 || !C)
			ParVec_AXPBY(alpha, z, gamma, x, y);
		else
		{
			C->Apply(x, y);
			ParVec_AXPBYPCZ(alpha, z, beta, y, gamma, x, y);
		}
	}
}

void Par_Compound_Operator::Apply(const Par_MultiVector& X, const Par_MultiVector& Y) const
{
	int m = X.num_vectors;
	if (alpha == 0 || !B)
	{
		if (beta == 0 || !C)
			for (int j = 0; j < m; ++j)
				ParVec_AXPBY(0.0, Y(j), gamma, X(j), Y(j));
		else
		{
			C->Apply(X, Y);
			for (int j = 0; j < m; ++j)
				ParVec_AXPBY(beta, Y(j), gamma, X(j), Y(j));
		}
	}
	else
	{
		Par_MultiVector Z(comm);

		if (A)
		{
			Par_MultiVector W(comm);
			W.Allocate(B->OutSize(), m);
			Z.Allocate(A->OutSize(), m);
			B->Apply(X, W);
			A->Apply(W, Z);
		}
		else
		{
			Z.Allocate(B->OutSize(), m);
			B->Apply(X, Z);
		}
		
		if (beta == 0 || !C)
		{
			for (int j = 0; j < m; ++j)
				ParVec_AXPBY(alpha, Z(j), gamma, X(j), Y(j));
		}
		else
		{
			C->Apply(X, Y);
			for (int j = 0; j < m; ++j)
				ParVec_AXPBYPCZ(alpha, Z(j), beta, Y(j), gamma, X(j), Y(j));
		}
	}
}

}